#include <iostream>
#include <cmath>
#include "alignment.h"
#include "MbMatrix.h"
#include "model.h"
#include "settings.h"
#include "tree.h"
#include "util.h"
#include "timer.h"

using namespace std;



Model::Model(Alignment *ap, MbRandom *rp, Settings *sp) {
    
	alignmentPtr = ap;
	ranPtr       = rp;
	settingsPtr  = sp;
	
	/* set up the conditional likelihoods */
	int numTaxa      = alignmentPtr->getNumTaxa();
    numNodes     = 2 * numTaxa - 2;
	int numChar      = alignmentPtr->getNumChar();
	int numStates    = 4;
	int condLikeSize = numNodes * numChar * numStates;
	cls = new double[condLikeSize];
	if ( !cls )
    {
		cerr << "ERROR: Problem allocating conditional likelihoods" << endl;
		exit(1);
    }
	for (int i=0; i<condLikeSize; i++)
		cls[i] = 0.0;
    
	/* allocate a matrix of pointers so we can quickly position ourselves in
     the vector of conditional likelihoods (cls) */
	clsPtr = new double**[numNodes];
	clsPtr[0] = new double*[numNodes * numChar];
	for (int i=1; i<numNodes; i++)
		clsPtr[i] = clsPtr[i-1] + numChar;
	for (int i=0; i<numNodes; i++)
		for (int j=0; j<numChar; j++)
			clsPtr[i][j] = &cls[i * numChar * numStates + j * numStates];
    
	/* initialize the conditional likelihoods at the tips of the tree, which
     are indexed 0, 1, ..., numTaxa-1 */
	for (int i=0; i<numTaxa; i++)
    {
		for (int j=0; j<numChar; j++)
        {
			double *clp = clsPtr[i][j];
			int nucCode = alignmentPtr->getNucleotide(i, j);
			int nucs[4];
			alignmentPtr->getPossibleNucs(nucCode, nucs);
			for (int k=0; k<numStates; k++)
            {
				if ( nucs[k] == 1 )
					clp[k] = 1.0;
            }
        }
    }
	
	/* make a tree */
	string treeStr = getLineFromFile( settingsPtr->getTreeFileName(), 0 ); // read the tree from the file
	tree = new Tree(alignmentPtr, ranPtr, treeStr);
	tree->print();
	
	/* set up transition probability matrices according to the Jukes-Cantor (1969) model of
     DNA substitution */
	P = new MbMatrix<double>[numNodes];
	for (int i=0; i<numNodes; i++)
		P[i] = MbMatrix<double>(4, 4);
	for (int n=0; n<numNodes; n++)
    {
		Node *p = tree->getDownPassNode( n );
		if (p->getAnc() != NULL)
        {
			int idx = p->getIndex();
			double v = p->getV();
			double x = exp(-(4.0/3.0)*v);
			for (int i=0; i<4; i++)
            {
				for (int j=0; j<4; j++)
                {
					if (i == j)
						P[idx][i][j] = 0.25 + 0.75 * x;
					else
						P[idx][i][j] = 0.25 - 0.25 * x;
                }
            }
        }
    }
	
}

Model::~Model(void) {
    
	delete [] cls;
	delete [] clsPtr[0];
	delete [] clsPtr;
	delete tree;
	
}

double Model::calculateLikelihood(void) {
    
    /* Calculate conditional likelihoods down the tree. Note that
     we have already calculated the transition probabilities in
     the constructor. */
    double *lnScaler = new double[alignmentPtr->getNumChar()];
    Timer likelihood;
    initialize_timer(&likelihood);
    start_timer(&likelihood);
    for (int c=0; c<alignmentPtr->getNumChar(); c++)
        lnScaler[c] = 0.0;
    
    for (int n=0; n<numNodes; n++)
    {
        Node *p = tree->getDownPassNode( n );
        if (p->getLft() != NULL && p->getRht() != NULL && p->getAnc() != NULL)
        {
            /* deal with a two-way split */
            if (p->getAnc()->getAnc() != NULL)
            {
                int lftIdx = p->getLft()->getIndex();
                int rhtIdx = p->getRht()->getIndex();
                int idx    = p->getIndex();
                double *clL = clsPtr[lftIdx][0];
                double *clR = clsPtr[rhtIdx][0];
                double *clP = clsPtr[idx   ][0];
                for (int c=0; c<alignmentPtr->getNumChar(); c++)
                {
                    double maxCl = 0.0;
                    for (int i=0; i<4; i++)
                    {
                        double sumL = 0.0, sumR = 0.0;
                        for (int j=0; j<4; j++)
                        {
                            sumL += clL[j] * P[lftIdx][i][j];
                            sumR += clR[j] * P[rhtIdx][i][j];
                        }
                        clP[i] = sumL * sumR;
                        if (clP[i] > maxCl)
                            maxCl = clP[i];
                    }
                    lnScaler[c] += log(maxCl);
                    for (int i=0; i<4; i++)
                        clP[i] /= maxCl;
                    clL += 4;
                    clR += 4;
                    clP += 4;
                }
            }
            else
            {
                /* deal with the node to the left of the root node */
                int lftIdx = p->getLft()->getIndex();
                int rhtIdx = p->getRht()->getIndex();
                int ancIdx = p->getAnc()->getIndex();
                int idx    = p->getIndex();
                double *clL = clsPtr[lftIdx][0];
                double *clR = clsPtr[rhtIdx][0];
                double *clA = clsPtr[ancIdx][0];
                double *clP = clsPtr[idx   ][0];
                for (int c=0; c<alignmentPtr->getNumChar(); c++)
                {
                    double maxCl = 0.0;
                    for (int i=0; i<4; i++)
                    {
                        double sumL = 0.0, sumR = 0.0, sumA = 0.0;
                        for (int j=0; j<4; j++)
                        {
                            sumL += clL[j] * P[lftIdx][i][j];
                            sumR += clR[j] * P[rhtIdx][i][j];
                            sumA += clA[j] * P[idx   ][i][j];
                        }
                        clP[i] = sumL * sumR * sumA;
                        if (clP[i] > maxCl)
                            maxCl = clP[i];
                    }
                    lnScaler[c] += log(maxCl);
                    for (int i=0; i<4; i++)
                        clP[i] /= maxCl;
                    clL += 4;
                    clR += 4;
                    clA += 4;
                    clP += 4;
                }
            }
        }
    }
    
	/* calculate likelihood */
	Node *p = tree->getRoot()->getLft();
	double *clP = clsPtr[p->getIndex()][0];
	double lnL = 0.0;
	for (int c=0; c<alignmentPtr->getNumChar(); c++)
    {
		double like = 0.0;
		for (int i=0; i<4; i++)
			like += clP[i] * 0.25;
		//cout << c << " -- " << log(like) << " " << alignmentPtr->getNumOfPattern(c) << endl;
		lnL += lnScaler[c] + log(like) * alignmentPtr->getNumOfPattern(c);
		clP += 4;
    }
	delete [] lnScaler;
    stop_timer(&likelihood);
    fprintf(stderr,"%f seconds\n",timer_duration(likelihood));
	
	return lnL;
    
}

